
from starlette.websockets import WebSocket

from wintersweet.framework.fastapi.request import APIRouter, Request
from wintersweet.asyncs.pool.mysql import mysql_pool_manager, atomic
from model import user_sa
from wintersweet.framework.fastapi.ws import WsChatConnectionManager

router = APIRouter()


@router.get('/')
async def test(request: Request):

    #
    # # 查单条
    # user: dict = await mysql_pool_manager.find(user_sa.select().where(user_sa.c.id == 1), pool_name='default')
    #
    # # 查多条
    # users: List[dict] = await mysql_pool_manager.select(user_sa.select().where(user_sa.c.id >= 1))
    #
    # # 修改
    # rows: int = await mysql_pool_manager.update(user_sa.update().where(user_sa.c.id == 1).values(username='dxx'))
    #
    # # 删除
    # await mysql_pool_manager.delete(user_sa.delete().where(user_sa.c.id == 1))
    #
    # # 新增
    # await mysql_pool_manager.insert(user_sa.insert().values(username='dxx'))
    # # 批量新增
    # await mysql_pool_manager.insert(user_sa.insert().values([{'username': 'dxx'}, {'username': 'pdd'}]))
    #
    # # 事务操作 方法1
    # await atomic_func()
    #
    # # 事务操作 方法2
    # async with mysql_pool_manager.get_client() as client:
    #     async with client.begin():
    #         user: dict = await mysql_pool_manager.find(user_sa.select().where(user_sa.c.id == 1))
    #         rows: int = await mysql_pool_manager.update(
    #             user_sa.update().where(user_sa.c.id == user['id']).values(username='dxx'))
    #
    # # 缓存操作
    # async with redis_pool_manager.get_client() as cache:
    #     await cache.set('test', 1)
    #     print(await cache.get('test'))

    # http请求
    # from wintersweet.asyncs import http
    #
    # resp = await http.get('https://www.baidu.com')
    # print(resp.text)
    # print(resp.status_code)
    # print(resp.headers)
    #
    # # 文件下载
    #
    # resp = await http.download('https://www.baidu.com', file_path='baidu.html')
    # print(resp.success)

    return request.client_ip


# 被atomic修饰的函数在执行之初就会消耗一个conn 并开启事务，在函数运行结束时自动commit
# 需要注意的是，请不要在大耗时函数上使用事务装饰器，这样会持续占用数据库连接，极端并发情况下可能导致性能问题，尽可能将数据库事务操作提炼到一个函数再进行修饰
@atomic(pool_name='default')
async def atomic_func():
    user: dict = await mysql_pool_manager.find(user_sa.select().where(user_sa.c.id == 1))
    rows: int = await mysql_pool_manager.update(user_sa.update().where(user_sa.c.id == user['id']).values(username='dxx'))


async def on_message(ws, message):
    ws.send_text(f' receive {message}')

manager = WsChatConnectionManager()


class MessageHandler:

    async def hello(self, ws, count):
        await ws.send_text('hello')
        pass

    @staticmethod
    async def on_message(self, ws, message):
        await ws.send_text(f'收到了"{message}"')

        pass


@router.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):

    key = id(ws)
    await manager.append(key, ws, MessageHandler())
    await manager.handle_message(key)